package com.shen.langchain4j.controller;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.lang.UUID;
import com.shen.langchain4j.service.FunctionAssistant;
import com.shen.langchain4j.service.InvoiceHandler;
import com.shen.langchain4j.service.WeatherService;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.service.tool.DefaultToolExecutor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.List;

@Slf4j
@RestController
@RequestMapping("/chatFunctionCalling")
public class ChatFunctionCallingController {
    @Autowired
    @Qualifier("qwen")
    private ChatModel chatModel;
    @Autowired
    private FunctionAssistant functionAssistant;
    @Autowired
    private WeatherService weatherService;

    /**
     * 低阶工具使用 构建工具规格说明(ToolSpecification) + 并手动执行(DefaultToolExecutor) + 大模型分析
     *
     * @return 大模型结合调用工具结果的回答
     */
    @RequestMapping("/lowLevelFunctionCallingChat")
    public String lowLevelFunctionCallingChat() {
        List<ToolSpecification> toolSpecifications = ToolSpecifications.toolSpecificationsFrom(new InvoiceHandler());
        List<ChatMessage> messageList = new ArrayList<>();
        UserMessage userMessage = UserMessage.from("镜片品牌推荐");
        ChatRequest chatRequest = ChatRequest.builder()
                .messages(userMessage)
                .toolSpecifications(toolSpecifications)
                .build();
        ChatResponse chatResponse = chatModel.chat(chatRequest);
        // 第一次提问如果是开具发票相关的，会把使用工具的参数返回，反之则直接给出回答
        AiMessage aiMessage = chatResponse.aiMessage();
        if (aiMessage.hasToolExecutionRequests()) {
            messageList.add(aiMessage);
            aiMessage.toolExecutionRequests().forEach(t -> {
                DefaultToolExecutor toolExecutor = new DefaultToolExecutor(new InvoiceHandler(), t);
                String result = toolExecutor.execute(t, UUID.randomUUID());
                log.info("工具执行后的结果为：{}", result);
                ToolExecutionResultMessage toolMessage = ToolExecutionResultMessage.from(t, result);
                messageList.add(toolMessage); // 将结果加入消息历史
            });
            AiMessage finalMessage = chatModel.chat(messageList).aiMessage();
            log.info("answer is {}", finalMessage);
            return "success :" + DateUtil.now() + "\n" + finalMessage.text();
        }
        return aiMessage.text();
    }

    /**
     * 高阶工具使用 使用@Tool注解指定方法，并在构建AI服务实例时添加，大模型会根据情况调用并自动执行工具
     *
     * @return 大模型结合调用工具结果的回答
     */
    @RequestMapping("/highLevelFunctionCallingChat")
    public String highLevelFunctionCallingChat() {
        //如果提问是开具发票相关的则使用工具，反之则不会
        String answer = functionAssistant.chat("根据以下信息开张发票," +
                "公司：北京深度求索人工智能基础技术研究有限公司" +
                "税号：192886685Z8LGVYSP2" +
                "金额：80000");
        log.info("answer is {}", answer);
        return "success :" + DateUtil.now() + "\n" + answer;
    }


    @RequestMapping("/queryCityWeather")
    public String queryCityWeather() throws Exception {
        //北京的location id 为101010100
        return weatherService.getWeather("101010100").toString();
    }
}
